from swarm import Swarm

from JiuTian import process_and_print_streaming_response, pretty_print_messages


def display(content):
    print(content)

#多轮对话公共函数
def run_multi_turn_conversation(
    openai_client,
    starting_agent,
    context_variables=None,
    debug=False
) -> None:
    # 创建 Swarm 客户端
    client = Swarm(openai_client)
    display("## 开启Swarm对话 ")

    # 初始化消息列表
    messages = []
    agent = starting_agent  # 初始智能体

    while True:
        # 从用户获取输入
        user_input = input("User: ")
        if user_input.lower() in ["exit", "quit"]:
            display("### Conversation Ended")
            break

        # 将用户输入添加到消息列表中
        messages.append({"role": "user", "content": user_input})

        # 运行 Swarm 客户端，智能体处理消息
        response = client.run(
            agent=agent,
            messages=messages,
            context_variables=context_variables or {},
            debug=debug,
        )

        # 使用 display(Markdown) 打印用户消息和智能体回复
        for message in response.messages:
            if message['role'] == 'user':
                display(f"**User**: {message['content']}")
            elif message['role'] == 'assistant':
                display(f"**{message['sender']}**: {message['content']}")

        # 更新消息和当前的智能体
        messages.extend(response.messages)
        agent = response.agent


def run_demo_loop(
        openai_client,
        starting_agent,
        context_variables=None,
        stream=False,
        debug=False) -> None:
    client = Swarm(openai_client)
    print("Starting Swarm CLI 🐝")
    print('Type "exit" or "quit" to leave the chat.')

    messages = []
    agent = starting_agent

    # 主循环，用户可以持续与智能体对话
    while True:
        user_input = input("\033[90mUser\033[0m: ").strip()  # 读取用户输入并去除首尾空格

        # 检查用户是否输入了退出关键词
        if user_input.lower() in {"exit", "quit"}:
            print("Exiting chat. Goodbye!")
            break  # 退出循环，结束聊天

        messages.append({"role": "user", "content": user_input})  # 将用户输入添加到消息列表

        # 运行 Swarm 客户端，与智能体交互
        response = client.run(
            agent=agent,
            messages=messages,
            context_variables=context_variables or {},
            stream=stream,
            debug=debug,
        )

        if stream:
            # 如果启用了流式处理，调用流处理函数
            response = process_and_print_streaming_response(response)
        else:
            # 否则直接打印消息
            pretty_print_messages(response.messages)

        # 更新消息和当前智能体
        messages.extend(response.messages)
        agent = response.agent